from datetime import datetime, timedelta,date
import pandas as pd
from pathlib import Path

# trading_dates: 普通交易日
trading_dates = pd.to_datetime(pd.read_csv(Path(__file__).parent / 'trade_dates.csv').date).dt.date

# holidays: 非交易日
all_dates = pd.date_range('2000-01-01', datetime.today() - timedelta(days=1), freq='D')
holidays = pd.Series(all_dates.difference(trading_dates).date,name='date')

# night_trading_dates: 有夜盘的交易日
df = pd.DataFrame(trading_dates)
df['datetime'] = pd.to_datetime(df.date)
df['weekday'] = df.datetime.dt.weekday + 1
df['date_diff'] = (df.date.diff().dt.days).shift(-1)
df = df[ (((df.weekday==5) & (df.date_diff == 3)) | (df.date_diff==1))  ].dropna()
df = df[(df.date<date(2020,2,3)) | (df.date>date(2020,5,1))] # 去除新冠疫情期间无夜盘的交易日
night_trading_dates = df.date

def parse_datetime(datetime_input: int | str | datetime | date) -> str:
    """将不同的时间格式转化为统一格式的字符串
    Returns:
        str: 'YYYY-mm-dd HH:MM:SS'
    """
    if datetime_input is None:
        raise Exception("datetime_input 日期输入不可为None")
    datetime_dt: pd.Timestamp = pd.to_datetime(str(datetime_input))
    return str(datetime_dt)

def to_datetime(datetime_input: int | str | datetime | date) -> datetime:
    """将不同的时间格式转化为datetime.datetime对象
    Returns:
        datetime.datetime
    """
    return pd.to_datetime(parse_datetime(datetime_input))

def get_trading_dates(start_date: any, end_date: any):
    left_idx = trading_dates.searchsorted(pd.to_datetime(str(start_date)).date(), side='left')
    right_idx = trading_dates.searchsorted(pd.to_datetime(str(end_date)).date(), side='right')
    return trading_dates[left_idx:right_idx].to_list()

def get_next_trading_date(date: any, n = 1):
    idx = trading_dates.searchsorted(pd.to_datetime(str(date)).date(),side = 'right')
    return trading_dates[idx + n - 1]

def get_previous_trading_date(date: any, n = 1):
    idx = trading_dates.searchsorted(pd.to_datetime(str(date)).date(),side = 'left')
    return trading_dates[idx - n]

def is_trading_date(date: any):
    return pd.to_datetime(parse_datetime(date)).date() in trading_dates.values

def has_night_trading(date: any):
    return pd.to_datetime(parse_datetime(date)).date() in night_trading_dates.values

def get_holidays(start_date: any, end_date: any):
    return holidays[(holidays >= pd.to_datetime(parse_datetime(start_date)).date()) & (holidays <= pd.to_datetime(parse_datetime(end_date)).date())].to_list()

def check_trading_date(func):
    def wrapper(trading_date, *args, **kwargs):
        if not is_trading_date(trading_date):
            raise Exception(f"{trading_date} 不是交易日")
        else:
            return func(trading_date, *args, **kwargs)
    return wrapper